import os
import pdb
from typing_extensions import Literal

from mmseg.datasets import BaseSegDataset
from mgamdata.dataset.AbdomenCT_1K.meta import CLASS_INDEX_MAP

class AbdomenCT1K_2D(BaseSegDataset):
    def __init__(self, 
                 data_root, 
                 split:Literal["train", "val", "test"], 
                 debug:bool, 
                 *args, **kwargs
    ):
        super().__init__(data_root=data_root, *args, **kwargs)
        self.split = split
        self.debug = debug
        """
        如果需要的话
        init函数也可以进行修改
        """
    
    def load_data_list(self) -> list[dict]:
        """
        需要实现这个函数，
        它的返回值是一个包含众多字典的列表，
        每个字典代表一个样本，
        字典中应当包含如下的键：
            - img_path
            - seg_map_path
             (下面三行的值直接这么写就行)
            - label_map=self.label_map
            - reduce_zero_label=False,
            - seg_fields=[]
        
        在实现时，需要注意根据split来返回不重叠的数据子集。
        如果debug为true，则只返回少量的样本。
        """
        
        samples = []
        example_image = "/mnt/wsl/PHYSICALDRIVE4/mgam_datasets/AbdomenCT_1K/spacing2_2D_png/image/00001_z100.png"
        example_label = "/mnt/wsl/PHYSICALDRIVE4/mgam_datasets/AbdomenCT_1K/spacing2_2D_png/label/00001_z100.png"
        
        samples.append({
            "img_path": example_image,
            "seg_map_path": example_label,
            "label_map": CLASS_INDEX_MAP,
            "reduce_zero_label": False,
            "seg_fields": []
        })
        
        return samples if not self.debug else samples[:10]
